package xyz.hzxiao.utils;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * Rsa
 *
 * @author ：FuXx
 * @className : cn.miying.utils.RsaDemo
 * @since ：2024/9/9
 */
public class RsaDemo {

    public static final String ALGORITHM = "RSA";

    /**
     * 密钥文件输出路径
     */
    public static final String OUT_FILE_PATH = "E:\\temp\\";

    public static void main(String[] args) throws Exception {
//            createNewKeyPairAndWriteToFile();
        Key publicKey = readKeyByPemFile(OUT_FILE_PATH + "public_key.pem");
        Key privateKey = readKeyByPemFile(OUT_FILE_PATH + "private_key.pem");
        System.out.println("公钥: " + Base64.getEncoder().encodeToString(publicKey.getEncoded()));
        System.out.println("私钥: " + Base64.getEncoder().encodeToString(privateKey.getEncoded()));
        String originalData = "FuX是个大帅比";
        String encryptedData = encrypt(privateKey, originalData);
        String decryptedData = decrypt(publicKey, encryptedData);

        System.out.println("加密后的数据: " + encryptedData);
        System.out.println("解密后的数据: " + decryptedData);
    }

    /**
     * 解密方法
     *
     * @param privateKey    私钥
     * @param encryptedData 要解密的数据
     * @return 解密后的数据
     */
    public static String decrypt(Key privateKey, String encryptedData) throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] decryptedBytes = cipher.doFinal(Base64.getDecoder().decode(encryptedData));
        return new String(decryptedBytes);
    }

    /**
     * 加密方法
     *
     * @param publicKey 公钥
     * @param data      要加密的数据
     * @return 加密后的数据
     */
    public static String encrypt(Key publicKey, String data) throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        Cipher cipher = Cipher.getInstance(ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        byte[] bytes = cipher.doFinal(data.getBytes());
        return Base64.getEncoder().encodeToString(bytes);
    }

    /**
     * 生成RSA密钥对,长度为2048
     *
     * @return 密钥对
     * @throws NoSuchAlgorithmException 算法异常
     */
    public static KeyPair generateRsaKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM);
        // 可以指定密钥长度，如2048位
        keyGen.initialize(2048);
        return keyGen.generateKeyPair();
    }

    /**
     * 创建新的密钥对并写入文件
     *
     * @throws NoSuchAlgorithmException 密钥对生成异常
     * @throws IOException              文件写入异常
     */
    public static void createNewKeyPairAndWriteToFile() throws NoSuchAlgorithmException, IOException {
        KeyPair keyPair = generateRsaKeyPair();
        Key publicKey = keyPair.getPublic();
        Key privateKey = keyPair.getPrivate();
        writeKeyToPemFile(publicKey, OUT_FILE_PATH + "public_key.pem");
        writeKeyToPemFile(privateKey, OUT_FILE_PATH + "private_key.pem");
    }

    /**
     * 将Key以PEM格式写入文件的方法,兼容公私钥
     *
     * @param key      密钥
     * @param filePath 输出文件路径
     * @throws IOException 文件写入异常
     */
    public static void writeKeyToPemFile(Key key, String filePath) throws IOException {
        byte[] publicKeyBytes = key.getEncoded();
        String base64PublicKey = Base64.getEncoder().encodeToString(publicKeyBytes);
        String pemFormatKey = "-----BEGIN " + (key instanceof PrivateKey ? "PRIVATE" : "PUBLIC") + " KEY-----\n" +
                base64PublicKey + "\n-----END " + (key instanceof PrivateKey ? "PRIVATE" : "PUBLIC") + " KEY-----";
        Files.write(Paths.get(filePath), pemFormatKey.getBytes(), StandardOpenOption.CREATE);
    }

    /**
     * 从PEM格式文件中读取Key的方法,兼容公私钥
     *
     * @param filePath PEM格式文件路径
     * @return 密钥
     * @throws IOException 文件读取异常
     */
    public static Key readKeyByPemFile(String filePath) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
        String pemContent = new String(Files.readAllBytes(Paths.get(filePath)));
        String keyContent = pemContent.replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s+", "");
        byte[] keyBytes = Base64.getDecoder().decode(keyContent);
        KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
        return pemContent.contains("PRIVATE") ? keyFactory.generatePrivate(new PKCS8EncodedKeySpec(keyBytes)) : keyFactory.generatePublic(new X509EncodedKeySpec(keyBytes));
    }
}
